#!/usr/bin/python3
# -*- coding: utf-8 -*-
import serial
import serial.tools.list_ports

import rospy
from geometry_msgs.msg import Transform,Pose,Point32
from cv_bridge import CvBridge, CvBridgeError

import numpy as np

from bs_img.serial_pack import *

from bs_img.bs_cfg_real_5m import *
 
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import pandas as pd


class get_final_res():
    def __init__(self) -> None:
        self.lf_drogue_res_topic = "/bs_debuger/result/lf_drogue"
        self.sf_drogue_res_topic = "/bs_debuger/result/sf_drogue"
        self.wa_drogue_res_topic = "/bs_debuger/result/wa_drogue"
        self.lf_plane_res_topic = "/bs_debuger/result/lf_plane"
        self.sf_plane_res_topic = "/bs_debuger/result/sf_plane"

        self.lf_drogue_res = [np.nan]*3
        self.sf_drogue_res = [np.nan]*3
        self.wa_drogue_res = [np.nan]*3
        self.lf_plane_res = [np.nan]*3
        self.sf_plane_res = [np.nan]*3
        # weight coefficient for the results of lf camera in fusion
        self.coeff = camK_lf[0][0]/(camK_lf[0][0]+camK_sf[0][0])
        self.sleep_interval = 0.0005


        #final result
        self.plane_res = np.array([np.nan, np.nan, np.nan])
        self.drogue_res = np.array([np.nan, np.nan, np.nan])

        self.lf_drogue_res_sub = rospy.Subscriber(self.lf_drogue_res_topic, Transform, self.lf_drogue_res_cb)
        self.sf_drogue_res_sub = rospy.Subscriber(self.sf_drogue_res_topic, Transform, self.sf_drogue_res_cb)
        self.wa_drogue_res_sub = rospy.Subscriber(self.wa_drogue_res_topic, Transform, self.wa_drogue_res_cb)
        self.lf_plane_res_sub = rospy.Subscriber(self.lf_plane_res_topic, Transform, self.lf_plane_res_cb)
        self.sf_plane_res_sub = rospy.Subscriber(self.sf_plane_res_topic, Transform, self.sf_plane_res_cb)

        # publish results
        self.drogue_res_pub = rospy.Publisher("/bs_debuger/result/drogue", Transform , queue_size=1)
        self.plane_res_pub = rospy.Publisher("/bs_debuger/result/plane", Transform , queue_size=1)




    def lf_drogue_res_cb(self,transform_msg):
        self.lf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        
    def sf_drogue_res_cb(self,transform_msg):
        self.sf_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        lf_drouge_is_visiable = True if not(np.all(np.isnan(self.lf_drogue_res))) else False
        sf_drouge_is_visiable = True if not(np.all(np.isnan(self.sf_drogue_res))) else False
        wa_drouge_is_visiable = True if not(np.all(np.isnan(self.wa_drogue_res))) else False
        if wa_drouge_is_visiable:
            self.drogue_res = self.wa_drogue_res
            drogue_res_frame = serial_frame("drogue",True,round(self.drogue_res[2]*100),round(self.drogue_res[0]*100),round(self.drogue_res[1]))        
        elif lf_drouge_is_visiable and sf_drouge_is_visiable:
            self.drogue_res = np.array(self.lf_drogue_res)*self.coeff + np.array(self.sf_drogue_res)*(1-self.coeff)
            drogue_res_frame = serial_frame("drogue",True,round(self.drogue_res[2]*100),round(self.drogue_res[0]*100),round(self.drogue_res[1]))      
        elif not(lf_drouge_is_visiable) and sf_drouge_is_visiable:
            self.drogue_res = self.sf_drogue_res
            drogue_res_frame = serial_frame("drogue",True,round(self.drogue_res[2]*100),round(self.drogue_res[0]*100),round(self.drogue_res[1]))        
        elif not(sf_drouge_is_visiable) and lf_drouge_is_visiable:
            self.drogue_res = self.lf_drogue_res
            drogue_res_frame = serial_frame("drogue",True,round(self.drogue_res[2]*100),round(self.drogue_res[0]*100),round(self.drogue_res[1]))    
        else:
            drogue_res_frame = serial_frame("drogue", False, round(0),round(0),round(0))

        uart.write(drogue_res_frame)
        time.sleep(self.sleep_interval)

        drogue_msg = Transform()
        drogue_msg.translation.x = self.drogue_res[0]
        drogue_msg.translation.y = self.drogue_res[1]
        drogue_msg.translation.z = self.drogue_res[2]
        self.drogue_res_pub.publish(drogue_msg)

    def wa_drogue_res_cb(self,transform_msg):
        self.wa_drogue_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]


    def lf_plane_res_cb(self,transform_msg):
        self.lf_plane_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        
    def sf_plane_res_cb(self,transform_msg):
        self.sf_plane_res = [transform_msg.translation.x, transform_msg.translation.y, transform_msg.translation.z]
        sf_plane_is_visiable = True if not(np.all(np.isnan(self.sf_plane_res))) else False
        lf_plane_is_visiable = True if not(np.all(np.isnan(self.lf_plane_res))) else False
        if not(sf_plane_is_visiable) and not(lf_plane_is_visiable):
            plane_res_frame = serial_frame("plane",False,round(0),round(0),round(0))
        elif sf_plane_is_visiable and not(lf_plane_is_visiable):
            self.plane_res = self.sf_plane_res
            plane_res_frame = serial_frame("plane",True,round((self.plane_res[2]-0.6)*100),round(self.plane_res[0]*100),round((self.plane_res[1]-0.13)*100))    
        elif lf_plane_is_visiable and not(sf_plane_is_visiable):
            self.plane_res = self.sf_plane_res
            plane_res_frame = serial_frame("plane",True,round((self.plane_res[2]-0.6)*100),round(self.plane_res[0]*100),round((self.plane_res[1]-0.13)*100))  
        else:
            self.plane_res = np.array(self.lf_plane_res)*self.coeff + np.array(self.sf_plane_res)*(1 - self.coeff)
            plane_res_frame = serial_frame("plane",True,round((self.plane_res[2]-0.6)*100),round(self.plane_res[0]*100),round((self.plane_res[1]-0.13)*100))  
        
        # need filtering

        uart.write(plane_res_frame)
        time.sleep(self.sleep_interval)

        plane_msg = Transform()
        plane_msg.translation.x = self.plane_res[0]
        plane_msg.translation.y = self.plane_res[1]
        plane_msg.translation.z = self.plane_res[2]
        self.plane_res_pub.publish(plane_msg)
        
        



if __name__ == "__main__":
    rospy.init_node('bs_solver_node', anonymous=True)

    port, baud =  '/dev/ttyUSB0', 115200
    uart = serial.Serial(port,baud,parity=serial.PARITY_NONE)

    final_res = get_final_res()

    rospy.spin()